#include "hdcameracalibration.h"
#include "hdshareddata.h"
#include <QImage>

//Width and height are just two dimensions. Orientation doesn't make any differences.
#define GRID_WIDTH 6
#define GRID_HEIGHT 7
#define GRID_BOX_DIMENSION 30.0f
#define GRID_CORNER_COUNT GRID_HEIGHT*GRID_WIDTH

#define SAVE_CALIBRATION_IMAGES true
#define SAVING_PATH "C:\\joe_calib4\\"
#define CAM_ZERO_NAME "left"
#define CAM_ONE_NAME "right"

#define CALIBRATION_IMAGE_COUNT 10.0f

bool bCalibrationFinished;
CvMat* _matUnproject;
bool bViewReady = false;

IplImage* tempZero;
IplImage* tempOne;

HDCameraCalibration::HDCameraCalibration(QObject *parent)
: QThread(parent)
{
	_arrCamera_Intrinsic = new CvMat*[NOOFCAMS];
	_arrCamera_Distortion = new CvMat*[NOOFCAMS];
	_UndistortMapX = new CvMat*[NOOFCAMS];
	_UndistortMapY = new CvMat*[NOOFCAMS];
	_RectifyMapX = new CvMat*[NOOFCAMS];
	_RectifyMapY =new CvMat*[NOOFCAMS];

	for (int i=0;i<NOOFCAMS;i++)
	{
		_arrCamera_Intrinsic[i]= cvCreateMat(3,3,CV_32FC1);
		_arrCamera_Distortion[i] = cvCreateMat(1,5,CV_32FC1);

		_UndistortMapX[i] = cvCreateMat(FRAME_HEIGHT,FRAME_WIDTH,CV_32FC1);
		_UndistortMapY[i] = cvCreateMat(FRAME_HEIGHT,FRAME_WIDTH,CV_32FC1);

		_RectifyMapX[i] = cvCreateMat(FRAME_HEIGHT,FRAME_WIDTH,CV_32FC1);
		_RectifyMapY[i] = cvCreateMat(FRAME_HEIGHT,FRAME_WIDTH,CV_32FC1);
	}

	_matCalibtationR = cvCreateMat(3,3,CV_64FC1);
	_matCalibrationT = cvCreateMat(3,1,CV_64FC1);
	_matCalibrationF = cvCreateMat(3,3,CV_64FC1);
	_matCalibrationE = cvCreateMat(3,3,CV_64FC1);

	_CheckboardSize = cvSize(GRID_WIDTH,GRID_HEIGHT);

	_ObjectPoints = cvCreateMat(1,CALIBRATION_IMAGE_COUNT*_CheckboardSize.height*_CheckboardSize.width,CV_32FC3);
	_ZeroImagesPoints = cvCreateMat(1,CALIBRATION_IMAGE_COUNT*_CheckboardSize.height*_CheckboardSize.width,CV_32FC2);
	_OneImagesPoints = cvCreateMat(1,CALIBRATION_IMAGE_COUNT*_CheckboardSize.height*_CheckboardSize.width,CV_32FC2);
	_PointCounts = cvCreateMat(1,CALIBRATION_IMAGE_COUNT,CV_32S);

	_matUnproject = cvCreateMat(4,4,CV_32FC1);

	_fEtalonUnit = GRID_BOX_DIMENSION;
	bCalibrationFinished = false;
	_bUndistortionPossible = false;
	_iCapturedFrames =0;

	tempZero = cvCreateImage(cvSize(FRAME_WIDTH, FRAME_HEIGHT), 8, 3);
	tempOne = cvCreateImage(cvSize(FRAME_WIDTH, FRAME_HEIGHT), 8, 3);

	InitObjectPoints();
	//Start the thread.
	start();
}

HDCameraCalibration::~HDCameraCalibration()
{
	for (int i=0;i<NOOFCAMS;i++)
	{
		cvReleaseMat(&_arrCamera_Intrinsic[i]);
		cvReleaseMat(&_arrCamera_Distortion[i]);

		cvReleaseMat(&_UndistortMapX[i]);
		cvReleaseMat(&_UndistortMapY[i]);

		cvReleaseMat(&_RectifyMapX[i]);
		cvReleaseMat(&_RectifyMapY[i]);
	}

	delete [] _arrCamera_Intrinsic;
	delete [] _arrCamera_Distortion;
	delete [] _UndistortMapX;
	delete [] _UndistortMapY;
	delete [] _RectifyMapX;
	delete []_RectifyMapY;

	cvReleaseMat(&_matCalibtationR);
	cvReleaseMat(&_matCalibrationT);
	cvReleaseMat(&_matCalibrationF);
	cvReleaseMat(&_matCalibrationE);

	cvReleaseMat(&_ObjectPoints);
	cvReleaseMat(&_ZeroImagesPoints);
	cvReleaseMat(&_OneImagesPoints);
	cvReleaseMat(&_PointCounts);

	cvReleaseImage(&tempZero);
	cvReleaseImage(&tempOne);
}

void HDCameraCalibration::InitObjectPoints()
{
	int pointIndex = 0;
	float pointX=0;
	float pointY=0;
	float pointZ=0;
	for(int i=0;i<CALIBRATION_IMAGE_COUNT;i++)
	{
		pointY=0;
		for (int j=0;j<_CheckboardSize.height;j++)
		{
			pointX=0;
			for (int k=0;k<_CheckboardSize.width;k++)
			{
				//pointIndex = (i*42 + j*7 + k *6)*3; 
				_ObjectPoints->data.fl[pointIndex] = pointX;
				_ObjectPoints->data.fl[pointIndex+1] = pointY;
				_ObjectPoints->data.fl[pointIndex+2] = pointZ;

				pointX+=_fEtalonUnit;
				pointIndex+=3;
			}		
			pointY+=_fEtalonUnit;

		}
	}
}

void HDCameraCalibration::run()
{
	forever
	{
		// Wait for all cameras to finish capturing.
		for (int i = 0; i < camerasCount; i++)
		{
			captureFinished[i]->acquire();
		}
		
		bViewReady = true;

		// Processing the images that has been captured.
		if(camerasCount>=2)
		{
			if(bCalibrationFinished)
			{
				RectifyingImages();
			}
			else if(_iCapturedFrames==CALIBRATION_IMAGE_COUNT)
			{
					Calibrating();
			}
			else if(bCalibrate)
				PreparingImages();
		}

		// Signal the Tracking/Viewing for each camera.
		calibrationFinished->release(camerasCount);

		bViewReady = false;
	}
}

void HDCameraCalibration::PreparingImages()
{

	IplImage* imgZero = cameras[0]->frame();
	IplImage* imgOne = cameras[1]->frame();

	IplImage* grayZero = cvCreateImage(cvGetSize(imgZero),IPL_DEPTH_8U,1);
	IplImage* grayOne = cvCreateImage(cvGetSize(imgOne),IPL_DEPTH_8U,1);

	cvCvtColor(imgZero,grayZero,CV_RGB2GRAY);
	cvCvtColor(imgOne,grayOne,CV_RGB2GRAY);

	CvPoint2D32f* cornersZero = new CvPoint2D32f[GRID_HEIGHT*GRID_WIDTH];
	CvPoint2D32f* cornersOne = new CvPoint2D32f[GRID_HEIGHT*GRID_WIDTH];

	//Search for chessboard
	int retZero= cvFindChessboardCorners(grayZero,_CheckboardSize,cornersZero);
	int retOne= cvFindChessboardCorners(grayOne,_CheckboardSize,cornersOne);

	//If both cameras find the checkboard
	if(retZero>0 && retOne>0)
	{
		emit PreparingImagesProgress(((float)_iCapturedFrames+1/CALIBRATION_IMAGE_COUNT)*100);

		//Draw the chessboard
		cvDrawChessboardCorners(imgZero,_CheckboardSize,cornersZero,GRID_WIDTH*GRID_HEIGHT,retZero);
		cvDrawChessboardCorners(imgOne,_CheckboardSize,cornersOne,GRID_WIDTH*GRID_HEIGHT,retOne);

		//Save +/- include the photo to calibration data.
		if(bSnap)
		{
			//Refine the discovered corners
			cvFindCornerSubPix(grayZero,cornersZero,GRID_CORNER_COUNT,
				cvSize(6, 6), cvSize(-1,-1) ,
				cvTermCriteria(CV_TERMCRIT_ITER+CV_TERMCRIT_EPS,30, 0.01));
			cvFindCornerSubPix(grayOne,cornersOne,GRID_CORNER_COUNT,
				cvSize(6, 6), cvSize(-1,-1) ,
				cvTermCriteria(CV_TERMCRIT_ITER+CV_TERMCRIT_EPS,30, 0.01));

			//Write the discovered corners to point lists
			int startIndex = _iCapturedFrames * GRID_CORNER_COUNT *2;
			int endIndex = (_iCapturedFrames+1) * GRID_CORNER_COUNT * 2;
			int cornerIndex =0;

			for (int i=startIndex;i<endIndex;i+=2)
			{
				_ZeroImagesPoints->data.fl[i] = cornersZero[cornerIndex].x;
				_ZeroImagesPoints->data.fl[i+1] = cornersZero[cornerIndex].y;

				_OneImagesPoints->data.fl[i] = cornersOne[cornerIndex].x;
				_OneImagesPoints->data.fl[i+1] = cornersOne[cornerIndex].y;

				cornerIndex++;
			}
			_PointCounts->data.i[_iCapturedFrames]=GRID_CORNER_COUNT;

			if(SAVE_CALIBRATION_IMAGES)
			{
				char number[3];
				itoa(_iCapturedFrames,number,10);
				char fileZero[50],fileOne[50];

				strcpy(fileZero,SAVING_PATH);
				strcpy(fileOne,SAVING_PATH);

				strcat(fileZero,CAM_ZERO_NAME);
				strcat(fileOne,CAM_ONE_NAME);

				strcat(fileZero,number);
				strcat(fileZero,".jpg");
				cvSaveImage(fileZero,grayZero);

				strcat(fileOne,number);
				strcat(fileOne,".jpg");
				cvSaveImage(fileOne,grayOne);
			}
			bSnap = false;
			_iCapturedFrames++;
		}
	}

	cvReleaseImage(&grayZero);
	cvReleaseImage(&grayOne);

	delete [] cornersZero;
	delete [] cornersOne;
}

void HDCameraCalibration::Calibrating()
{
	emit CalibrationStarted();

	cvCalibrateCamera2(_ObjectPoints,_ZeroImagesPoints,_PointCounts,cvSize(FRAME_WIDTH,FRAME_HEIGHT),_arrCamera_Intrinsic[0],_arrCamera_Distortion[0]);
	cvCalibrateCamera2(_ObjectPoints,_OneImagesPoints,_PointCounts,cvSize(FRAME_WIDTH,FRAME_HEIGHT),_arrCamera_Intrinsic[1],_arrCamera_Distortion[1]);
	cvStereoCalibrate(_ObjectPoints,_ZeroImagesPoints,_OneImagesPoints,_PointCounts,
		_arrCamera_Intrinsic[0],_arrCamera_Distortion[0],_arrCamera_Intrinsic[1],_arrCamera_Distortion[1],
		cvSize(FRAME_WIDTH,FRAME_HEIGHT), _matCalibtationR, _matCalibrationT, _matCalibrationE, _matCalibrationF,
		cvTermCriteria( CV_TERMCRIT_ITER+CV_TERMCRIT_EPS, 100, 1e-5),
		CV_CALIB_FIX_ASPECT_RATIO +CV_CALIB_ZERO_TANGENT_DIST +CV_CALIB_SAME_FOCAL_LENGTH );

	//Init Undistort Maps
	cvInitUndistortMap(_arrCamera_Intrinsic[0],_arrCamera_Distortion[0],_UndistortMapX[0],_UndistortMapY[0]);
	cvInitUndistortMap(_arrCamera_Intrinsic[1],_arrCamera_Distortion[1],_UndistortMapX[1],_UndistortMapY[1]);

	//Init Rectify Maps
	CvMat* Rl = cvCreateMat(3,3,CV_64FC1);
	CvMat* Rr = cvCreateMat(3,3,CV_64FC1);
	CvMat* Pl = cvCreateMat(3,4,CV_64FC1);
	CvMat* Pr = cvCreateMat(3,4,CV_64FC1);

	cvStereoRectify(_arrCamera_Intrinsic[0],_arrCamera_Intrinsic[1],_arrCamera_Distortion[0],_arrCamera_Distortion[1],
		cvSize(FRAME_WIDTH,FRAME_HEIGHT),_matCalibtationR,_matCalibrationT,Rl,Rr,Pl,Pr,_matUnproject);

	cvInitUndistortRectifyMap(_arrCamera_Intrinsic[0],_arrCamera_Distortion[0],Rl,Pl,_RectifyMapX[0],_RectifyMapY[0]);
	cvInitUndistortRectifyMap(_arrCamera_Intrinsic[1],_arrCamera_Distortion[1],Rr,Pr,_RectifyMapX[1],_RectifyMapY[1]);

	cvReleaseMat(&Rl);
	cvReleaseMat(&Rr);
	cvReleaseMat(&Pl);
	cvReleaseMat(&Pr);

	emit CalibrationFinished();
	bCalibrationFinished = true;
	_bUndistortionPossible=true;

	//TODO: Remove this automatic saving
	SaveRectificationMaps("c:\\zerox.hd","c:\\zeroy.hd","c:\\onex.hd","c:\\oney.hd","c:\\unproject.hd");
}

void HDCameraCalibration::RectifyingImages()
{
	IplImage* imgZero = cameras[0]->frame();
	IplImage* imgOne = cameras[1]->frame();

// 	IplImage* tempZero = cvCreateImage(cvGetSize(imgZero),imgZero->depth,imgZero->nChannels);
// 	IplImage* tempOne = cvCreateImage(cvGetSize(imgOne),imgOne->depth,imgOne->nChannels);

	if(bRectify)
	{
		cvRemap(imgZero,tempZero,_RectifyMapX[0],_RectifyMapY[0]);
		cvRemap(imgOne,tempOne,_RectifyMapX[1],_RectifyMapY[1]);

		cvCopy(tempZero,imgZero);
		cvCopy(tempOne,imgOne);
		
		//The following code shows the disparity maps instead of the rectified images. But it's sooo slow. 1 Frame per 4-5 seconds
		//CvSize size = cvGetSize(tempZero);
		//IplImage* grayZero = cvCreateImage(size,IPL_DEPTH_8U,1);
		//IplImage* grayOne = cvCreateImage(size,IPL_DEPTH_8U,1);

		//cvCvtColor(tempZero,grayZero,CV_RGB2GRAY);
		//cvCvtColor(tempOne,grayOne,CV_RGB2GRAY);

		//CvMat* disparity_left = cvCreateMat( size.height, size.width, CV_16S );
		//CvMat* disparity_right = cvCreateMat( size.height, size.width, CV_16S );
		//CvStereoGCState* state = cvCreateStereoGCState( 20, 2 );
		//cvFindStereoCorrespondenceGC( grayZero, grayOne,
		//	disparity_left, disparity_right, state, 0 );
		//cvReleaseStereoGCState( &state );

		//CvMat* disparity_left_visual = cvCreateMat( size.height, size.width, CV_8U );
		//cvConvertScale( disparity_left, disparity_left_visual, -16 );

		//CvMat* disparity_right_visual = cvCreateMat( size.height, size.width, CV_8U );
		//cvConvertScale( disparity_right, disparity_right_visual, 1 );

		//cvCvtColor(disparity_left_visual,imgZero,CV_GRAY2RGB);
		//cvCvtColor(disparity_right_visual,imgOne,CV_GRAY2RGB);

		//cvSaveImage("c:\\left.jpg",imgZero);
		//cvSaveImage("c:\\right.jpg",imgOne);
		//bRectify = false;

		//cvReleaseImage(&grayZero);
		//cvReleaseImage(&grayOne);
		//cvReleaseMat(&disparity_left);
		//cvReleaseMat(&disparity_right);
		//cvReleaseMat(&disparity_left_visual);
		//cvReleaseMat(&disparity_right_visual);
	}
	else if (_bUndistortionPossible)
	{
		cvRemap(imgZero,tempZero,_UndistortMapX[0],_UndistortMapY[0]);
		cvRemap(imgOne,tempOne,_UndistortMapX[1],_UndistortMapY[1]);
		cvCopy(tempZero,imgZero);
		cvCopy(tempOne,imgOne);
	}

	/*cameras[0]->frame = tempZero;
	cameras[1]->frame = tempOne;*/

// 	cvReleaseImage(&tempZero);
// 	cvReleaseImage(&tempOne);
}

bool HDCameraCalibration::isCalibrationFinished()
{
	return bCalibrationFinished;
}

bool HDCameraCalibration::SaveRectificationMaps(char* pathCamZeroX,char* pathCamZeroY, char* pathCamOneX, char* pathCamOneY, char* pathMatUnprojection)
{
	if(bCalibrationFinished)
	{
		cvSave(pathCamZeroX,_RectifyMapX[0]);
		cvSave(pathCamZeroY,_RectifyMapY[0]);
		cvSave(pathCamOneX,_RectifyMapX[1]);
		cvSave(pathCamOneY,_RectifyMapY[1]);
		cvSave(pathMatUnprojection,_matUnproject);
		return true;
	}
	else 
		return false;
}

bool HDCameraCalibration::LoadRectificationMaps(char* pathCamZeroX,char* pathCamZeroY, char* pathCamOneX, char* pathCamOneY, char* pathMatUnprojection)
{
	//TODO: Check memory leak. Should I delete the matrices before changing what they point to ?
	_RectifyMapX[0] = (CvMat*)cvLoad(pathCamZeroX);
	_RectifyMapY[0] = (CvMat*)cvLoad(pathCamZeroY);
	_RectifyMapX[1] = (CvMat*)cvLoad(pathCamOneX);
	_RectifyMapY[1] = (CvMat*)cvLoad(pathCamOneY);
	_matUnproject = (CvMat*) cvLoad(pathMatUnprojection);
	bCalibrationFinished = true;
	return true;
}